package wrr

import (
	"errors"
	"google.golang.org/grpc/balancer"
	"google.golang.org/grpc/balancer/base"
	"sync"
	"sync/atomic"
	"time"
)

const Name = "custom_weighted_round_robin"

var (
	errLimited  = errors.New("trigger current limiting")
	errBreaker  = errors.New("trigger circuit breaker")
	codeLimited = "1999"
	codeBreaker = "1998"
)

func newBuilder() balancer.Builder {
	return base.NewBalancerBuilder(Name, &PickerBuilder{}, base.Config{HealthCheck: true})
}

func init() {
	balancer.Register(newBuilder())
}

type PickerBuilder struct {
}

func (p *PickerBuilder) Build(info base.PickerBuildInfo) balancer.Picker {
	conns := make([]*weightConn, 0, len(info.ReadySCs))
	for sc, sci := range info.ReadySCs {
		md, _ := sci.Address.Metadata.(map[string]any)
		weightVal, _ := md["weight"]
		weight, _ := weightVal.(float64)
		co := &weightConn{
			SubConn:       sc,
			weight:        int(weight),
			currentWeight: int(weight),
		}
		co.available.Store(true)
		conns = append(conns, co)
	}
	return &Picker{
		conns: conns,
	}
}

type Picker struct {
	conns []*weightConn
	lock  sync.Mutex
}

func (p *Picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
	p.lock.Lock()
	defer p.lock.Unlock()
	if len(p.conns) == 0 {
		return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
	}
	// 总权重
	var total int
	var maxCC *weightConn
	for _, c := range p.conns {
		//job15
		if !c.available.Load() { //不可用
			continue
		}
		total += c.weight                                          //计算总的权重
		c.currentWeight = c.currentWeight + c.weight               //计算当前每个节点的权重
		if maxCC == nil || maxCC.currentWeight < c.currentWeight { //挑选权重最大的节点
			maxCC = c
		}
	}

	maxCC.currentWeight = maxCC.currentWeight - total //更新挑选出来节点的权重 避免连续选中

	return balancer.PickResult{
		SubConn: maxCC.SubConn,
		Done: func(info balancer.DoneInfo) {
			/*
				在负载均衡算法里面处理服务端返回的错误。你只需要关心两个错误，触发了限流和触发了熔断的错误。
				如果是触发了限流，那么调低节点的权重。
				如果是触发了熔断，那么将节点暂时挪出可用节点列表。
				在挪出可用节点列表之后，另外启动 goroutine，定时发送健康检查请求到服务端。
				如果服务端返回了正常响应，那么就挪回去可用节点列表。在挪回去可用节点列表的时候，要注意控制流量。
			*/
			//job15

			if info.Err == nil {
				//正确响应 提高初始权重
				maxCC.weight = max(maxCC.weight+1, 100)
			} else {
				codes := info.Trailer.Get("resp_code")
				var code = "200"
				if len(codes) > 0 {
					code = codes[0]
				}
				if code == codeLimited { //限流
					maxCC.weight = min(maxCC.weight-1, 1)
				} else if code == codeBreaker { //挪出可用列表
					maxCC.available.Store(false)
					go p.CheckConnAvailableStatus(maxCC)
				}
			}

			//job14 需要考虑 极限的情况 不允许降低到负数，不允许提高到某个极大的值
			//假设最低初始权重为1，最高初始权重为100
			//if info.Err != nil { //有错误 降低初始权重
			//	maxCC.weight = min(maxCC.weight-1, 1)
			//} else { //正确响应 提高初始权重
			//	maxCC.weight = max(maxCC.weight+1, 100)
			//}
			// 要在这里进一步调整weight/currentWeight
			// failover 要在这里做文章
			// 根据调用结果的具体错误信息进行容错
			// 1. 如果要是触发了限流了，
			// 1.1 你可以考虑直接挪走这个节点，后面再挪回来
			// 1.2 你可以考虑直接将 weight/currentWeight 调整到极低
			// 2. 触发了熔断呢？
			// 3. 降级呢？
		},
	}, nil

}

type weightConn struct {
	balancer.SubConn
	weight        int //初始权重
	currentWeight int //当前权重

	// 可以用来标记不可用
	available atomic.Bool
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}
func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}
func (p *Picker) CheckConnAvailableStatus(conn *weightConn) {
	healthCheckDuration, retryCnt := time.Minute, 0
	for {
		time.Sleep(healthCheckDuration)
		//发送健康检查到服务端 todo 不知道怎么发送
		flag := healthCheck(conn)
		if flag {
			conn.available.Store(true) //标记节点可用
			conn.weight = 1            //初始权重设置为最低 以便实现流量控制
			break
		}
		retryCnt++
		if retryCnt > 100 {
			//重试大于固定次数后如果节点还没起来就报警
			break
		}
	}
}

func healthCheck(conn balancer.SubConn) bool {
	return true
}
